package com.example.demo.service.impl;

import com.example.demo.entity.SseEmitterUTF8;
import com.example.demo.service.SseService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Service
@Slf4j
public class SseServiceImpl implements SseService {
    private static final Map<String, SseEmitterUTF8> SSE_CACHE = new ConcurrentHashMap<>();
    private static final Map<String, List<String>> KEY_MAP = new ConcurrentHashMap<>();
    @Resource
    private ApplicationEventPublisher eventPublisher;

    @Override
    public SseEmitterUTF8 getConn(String clientId, String uuid) {
        String key = clientId + "|" + uuid;
        final SseEmitterUTF8 sseEmitter = SSE_CACHE.get(key);

        if (sseEmitter != null) {
            return sseEmitter;
        } else {
            // 设置连接超时时间，需要配合配置项 spring.mvc.async.request-timeout: 600000 一起使用
            final SseEmitterUTF8 emitter = new SseEmitterUTF8(600_000L);
            // 注册超时回调，超时后触发
            emitter.onTimeout(() -> {
                log.info("连接已超时，正准备关闭，clientId = {}", key);
                SSE_CACHE.remove(key);
            });
            // 注册完成回调，调用 emitter.complete() 触发
            emitter.onCompletion(() -> {
                log.info("连接已关闭，正准备释放，clientId = {}", key);
                SSE_CACHE.remove(key);
                log.info("连接已释放，clientId = {}", key);
            });
            // 注册异常回调，调用 emitter.completeWithError() 触发
            emitter.onError(throwable -> {
                log.error("连接已异常，正准备关闭，clientId = {}", key, throwable);
                SSE_CACHE.remove(key);
            });

            SSE_CACHE.put(key, emitter);
            List<String> keyList = KEY_MAP.get(clientId);
            if (CollectionUtils.isEmpty(keyList)) {
                keyList = new ArrayList<>();
            }
            keyList.add(key);
            KEY_MAP.put(clientId, keyList);
            return emitter;
        }
    }

    /**
     * 模拟类似于 chatGPT 的流式推送回答
     *
     * @param clientId 客户端 id
     * @throws IOException 异常
     */
    @Override
    public void send(String clientId, String param) throws IOException {

        List<String> list = KEY_MAP.get(clientId);
        for (String key : list) {
            SseEmitterUTF8 emitter = SSE_CACHE.get(key);
            if (emitter != null) {
                // 推流内容到客户端
                emitter.send(param);
                // 结束推流
                // emitter.complete();
            }
        }
    }

    @Override
    public void closeConn(String clientId) {
        final SseEmitterUTF8 sseEmitter = SSE_CACHE.get(clientId);
        if (sseEmitter != null) {
            sseEmitter.complete();
        }
        SSE_CACHE.remove(clientId);
    }



}
